{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "bb425bf5",
   "metadata": {},
   "source": [
    "### Video Navigation with Qwen3-Omni"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b92d046f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ['VLLM_USE_V1'] = '0'\n",
    "os.environ['VLLM_WORKER_MULTIPROC_METHOD'] = 'spawn'\n",
    "os.environ[\"VLLM_LOGGING_LEVEL\"] = \"ERROR\"\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = \"0\"\n",
    "import torch\n",
    "import warnings\n",
    "import numpy as np\n",
    "\n",
    "warnings.filterwarnings('ignore')\n",
    "warnings.filterwarnings('ignore', category=DeprecationWarning)\n",
    "warnings.filterwarnings('ignore', category=FutureWarning)\n",
    "warnings.filterwarnings('ignore', category=UserWarning)\n",
    "\n",
    "from qwen_omni_utils import process_mm_info\n",
    "from transformers import Qwen3OmniMoeProcessor\n",
    "\n",
    "def _load_model_processor():\n",
    "    if USE_TRANSFORMERS:\n",
    "        from transformers import Qwen3OmniMoeForConditionalGeneration\n",
    "        if TRANSFORMERS_USE_FLASH_ATTN2:\n",
    "            model = Qwen3OmniMoeForConditionalGeneration.from_pretrained(MODEL_PATH,\n",
    "                                                                         dtype='auto',\n",
    "                                                                         attn_implementation='flash_attention_2',\n",
    "                                                                         device_map=\"auto\")\n",
    "        else:\n",
    "            model = Qwen3OmniMoeForConditionalGeneration.from_pretrained(MODEL_PATH, device_map=\"auto\", dtype='auto')\n",
    "    else:\n",
    "        from vllm import LLM\n",
    "        model = LLM(\n",
    "            model=MODEL_PATH, trust_remote_code=True, gpu_memory_utilization=0.95,\n",
    "            tensor_parallel_size=torch.cuda.device_count(),\n",
    "            limit_mm_per_prompt={'image': 1, 'video': 3, 'audio': 3},\n",
    "            max_num_seqs=1,\n",
    "            max_model_len=32768,\n",
    "            seed=1234,\n",
    "        )\n",
    "\n",
    "    processor = Qwen3OmniMoeProcessor.from_pretrained(MODEL_PATH)\n",
    "    return model, processor\n",
    "\n",
    "def run_model(model, processor, messages, return_audio, use_audio_in_video):\n",
    "    if USE_TRANSFORMERS:\n",
    "        text = processor.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)\n",
    "        audios, images, videos = process_mm_info(messages, use_audio_in_video=use_audio_in_video)\n",
    "        inputs = processor(text=text, audio=audios, images=images, videos=videos, return_tensors=\"pt\", padding=True, use_audio_in_video=use_audio_in_video)\n",
    "        inputs = inputs.to(model.device).to(model.dtype)\n",
    "        text_ids, audio = model.generate(**inputs, \n",
    "                                            thinker_return_dict_in_generate=True,\n",
    "                                            thinker_max_new_tokens=8192, \n",
    "                                            thinker_do_sample=False,\n",
    "                                            speaker=\"Ethan\", \n",
    "                                            use_audio_in_video=use_audio_in_video,\n",
    "                                            return_audio=return_audio)\n",
    "        response = processor.batch_decode(text_ids.sequences[:, inputs[\"input_ids\"].shape[1] :], skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]\n",
    "        if audio is not None:\n",
    "            audio = np.array(audio.reshape(-1).detach().cpu().numpy() * 32767).astype(np.int16)\n",
    "        return response, audio\n",
    "    else:\n",
    "        from vllm import SamplingParams\n",
    "        sampling_params = SamplingParams(temperature=1e-2, top_p=0.1, top_k=1, max_tokens=8192)\n",
    "        text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)\n",
    "        audios, images, videos = process_mm_info(messages, use_audio_in_video=use_audio_in_video)\n",
    "        inputs = {'prompt': text, 'multi_modal_data': {}, \"mm_processor_kwargs\": {\"use_audio_in_video\": use_audio_in_video}}\n",
    "        if images is not None: inputs['multi_modal_data']['image'] = images\n",
    "        if videos is not None: inputs['multi_modal_data']['video'] = videos\n",
    "        if audios is not None: inputs['multi_modal_data']['audio'] = audios\n",
    "        outputs = model.generate(inputs, sampling_params=sampling_params)\n",
    "        response = outputs[0].outputs[0].text\n",
    "        return response, None\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d37dcedc",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Unrecognized keys in `rope_scaling` for 'rope_type'='default': {'mrope_section', 'interleaved'}\n",
      "Unrecognized keys in `rope_scaling` for 'rope_type'='default': {'mrope_section', 'interleaved'}\n",
      "`torch_dtype` is deprecated! Use `dtype` instead!\n",
      "2025-09-14 20:56:01,784\tINFO util.py:154 -- Missing packages: ['ipywidgets']. Run `pip install -U ipywidgets`, then restart the notebook server for rich notebook output.\n",
      "You are attempting to use Flash Attention 2 without specifying a torch dtype. This might lead to unexpected behaviour\n",
      "Loading safetensors checkpoint shards:   0% Completed | 0/15 [00:00<?, ?it/s]\n",
      "Loading safetensors checkpoint shards:   7% Completed | 1/15 [00:00<00:09,  1.50it/s]\n",
      "Loading safetensors checkpoint shards:  20% Completed | 3/15 [00:01<00:06,  1.77it/s]\n",
      "Loading safetensors checkpoint shards:  27% Completed | 4/15 [00:02<00:08,  1.27it/s]\n",
      "Loading safetensors checkpoint shards:  33% Completed | 5/15 [00:03<00:05,  1.76it/s]\n",
      "Loading safetensors checkpoint shards:  40% Completed | 6/15 [00:04<00:06,  1.30it/s]\n",
      "Loading safetensors checkpoint shards:  47% Completed | 7/15 [00:05<00:07,  1.10it/s]\n",
      "Loading safetensors checkpoint shards:  53% Completed | 8/15 [00:06<00:06,  1.02it/s]\n",
      "Loading safetensors checkpoint shards:  60% Completed | 9/15 [00:07<00:06,  1.03s/it]\n",
      "Loading safetensors checkpoint shards:  67% Completed | 10/15 [00:08<00:05,  1.09s/it]\n",
      "Loading safetensors checkpoint shards:  73% Completed | 11/15 [00:10<00:04,  1.13s/it]\n",
      "Loading safetensors checkpoint shards:  80% Completed | 12/15 [00:11<00:03,  1.14s/it]\n",
      "Loading safetensors checkpoint shards:  87% Completed | 13/15 [00:12<00:02,  1.17s/it]\n",
      "Loading safetensors checkpoint shards:  93% Completed | 14/15 [00:13<00:01,  1.18s/it]\n",
      "Loading safetensors checkpoint shards: 100% Completed | 15/15 [00:14<00:00,  1.18s/it]\n",
      "Loading safetensors checkpoint shards: 100% Completed | 15/15 [00:14<00:00,  1.00it/s]\n",
      "\n",
      "The image processor of type `Qwen2VLImageProcessor` is now loaded as a fast processor by default, even if the model checkpoint was saved with a slow processor. This is a breaking change and may produce slightly different outputs. To continue using the slow processor, instantiate this class with `use_fast=False`. Note that this behavior will be extended to all models in a future release.\n",
      "Unused or unrecognized kwargs: images.\n",
      "Capturing CUDA graph shapes: 100%|██████████| 1/1 [00:00<00:00,  1.67it/s]\n"
     ]
    }
   ],
   "source": [
    "import librosa\n",
    "import audioread\n",
    "\n",
    "from IPython.display import Video\n",
    "from IPython.display import Image\n",
    "from IPython.display import Audio\n",
    "\n",
    "MODEL_PATH = \"Qwen/Qwen3-Omni-30B-A3B-Instruct\"\n",
    "# MODEL_PATH = \"Qwen/Qwen3-Omni-30B-A3B-Thinking\"\n",
    "\n",
    "USE_TRANSFORMERS = False\n",
    "TRANSFORMERS_USE_FLASH_ATTN2 = True\n",
    "\n",
    "model, processor = _load_model_processor()\n",
    "\n",
    "USE_AUDIO_IN_VIDEO = False\n",
    "RETURN_AUDIO = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b1330dc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<video src=\"https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen3-Omni/cookbook/video2.mp4\" controls  width=\"640\"  height=\"360\">\n",
       "      Your browser does not support the <code>video</code> element.\n",
       "    </video>"
      ],
      "text/plain": [
       "<IPython.core.display.Video object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Adding requests: 100%|██████████| 1/1 [00:00<00:00, 36.99it/s]\n",
      "Processed prompts: 100%|██████████| 1/1 [00:00<00:00,  9.21it/s, est. speed input: 5822.70 toks/s, output: 18.63 toks/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "left\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "USE_AUDIO_IN_VIDEO = False\n",
    "video_path = \"https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen3-Omni/cookbook/video2.mp4\"\n",
    "\n",
    "messages = [\n",
    "    {\n",
    "        \"role\": \"user\",\n",
    "        \"content\": [\n",
    "            {\"type\": \"video\", \"video\": video_path},\n",
    "            {\"type\": \"text\", \"text\": \"If I want to stop at the window. Which direction should I take?\"} \n",
    "        ]\n",
    "    }\n",
    "]\n",
    "\n",
    "display(Video(video_path, width=640, height=360))\n",
    "\n",
    "response, _ = run_model(model=model, messages=messages, processor=processor, return_audio=RETURN_AUDIO, use_audio_in_video=USE_AUDIO_IN_VIDEO)\n",
    "\n",
    "print(response)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "qwen3_omni_public",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
